﻿"""Baidu image classification helper for advanced_general API."""

import base64
import io
import json
import os
import re
import time
from typing import Callable, Dict, Iterable, List, Optional, Tuple
from urllib import error, parse, request

from PIL import Image

try:
    RESAMPLE = Image.Resampling.LANCZOS
except AttributeError:
    RESAMPLE = Image.LANCZOS



class BaiduAPIError(RuntimeError):
    """Raised when Baidu AI platform returns an error."""

    def __init__(self, message: str, payload: Optional[dict] = None) -> None:
        super().__init__(message)
        self.payload = payload or {}


class BaiduImageClassifier:
    """Simple client for Baidu advanced_general image classification API."""

    TOKEN_URL = "https://aip.baidubce.com/oauth/2.0/token"
    API_URL = "https://aip.baidubce.com/rest/2.0/image-classify/v2/advanced_general"
    INVALID_CHARS = re.compile(r'[<>:"/\\|?*]')
    MAX_PIXEL_SIDE = 4096
    TARGET_PIXEL_SIDE = 2000
    MAX_IMAGE_BYTES = 4 * 1024 * 1024
    JPEG_QUALITY_STEPS = (90, 85, 80, 75, 70, 60, 50, 40)

    def __init__(
        self,
        api_key: str,
        secret_key: str,
        min_score: float = 0.5,
        allow_unclassified: bool = True,
    ) -> None:
        self.api_key = api_key.strip()
        self.secret_key = secret_key.strip()
        self.min_score = min_score
        self.allow_unclassified = allow_unclassified
        self._access_token: Optional[str] = None
        self._token_expire_at: float = 0.0

    # ------------------------------------------------------------------
    def _ensure_access_token(self) -> str:
        if self._access_token and time.time() < self._token_expire_at - 60:
            return self._access_token

        if not self.api_key or not self.secret_key:
            raise BaiduAPIError("缺少 API Key 或 Secret Key，无法获取 access_token")

        data = parse.urlencode(
            {
                "grant_type": "client_credentials",
                "client_id": self.api_key,
                "client_secret": self.secret_key,
            }
        ).encode("utf-8")

        req = request.Request(self.TOKEN_URL, data=data)
        req.add_header("Content-Type", "application/x-www-form-urlencoded")

        try:
            with request.urlopen(req, timeout=10) as resp:
                payload = json.loads(resp.read().decode("utf-8"))
        except error.URLError as exc:
            raise BaiduAPIError(f"获取 access_token 失败: {exc}") from exc

        if "error" in payload or "error_description" in payload:
            raise BaiduAPIError(
                payload.get("error_description", "获取 access_token 失败"), payload
            )

        access_token = payload.get("access_token")
        if not access_token:
            raise BaiduAPIError("响应中缺少 access_token", payload)

        expires_in = int(payload.get("expires_in", 2592000))
        self._access_token = access_token
        self._token_expire_at = time.time() + max(expires_in - 30, 60)
        return access_token

    # ------------------------------------------------------------------
    def _post_image(self, image_bytes: bytes) -> dict:
        token = self._ensure_access_token()
        url = f"{self.API_URL}?access_token={token}"

        encoded_image = base64.b64encode(image_bytes).decode("utf-8")
        payload = parse.urlencode({"image": encoded_image}).encode("utf-8")

        req = request.Request(url, data=payload)
        req.add_header("Content-Type", "application/x-www-form-urlencoded")

        try:
            with request.urlopen(req, timeout=15) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except error.URLError as exc:
            raise BaiduAPIError(f"调用百度接口失败: {exc}") from exc

    def _prepare_image_bytes(self, image_path: str) -> bytes:
        file_size = os.path.getsize(image_path)

        try:
            with Image.open(image_path) as img:
                img.load()
                width, height = img.size
        except Exception as exc:
            if file_size <= self.MAX_IMAGE_BYTES:
                with open(image_path, "rb") as file_obj:
                    return file_obj.read()
            raise BaiduAPIError(f"无法打开图片：{exc}")

        max_side = max(width, height) if width and height else 0
        if file_size <= self.MAX_IMAGE_BYTES and max_side <= self.TARGET_PIXEL_SIDE:
            with open(image_path, "rb") as file_obj:
                return file_obj.read()

        try:
            with Image.open(image_path) as image:
                image.load()
                image = image.convert("RGB")
                width, height = image.size
                max_side = max(width, height) if width and height else 0
                resize_target = min(self.TARGET_PIXEL_SIDE, self.MAX_PIXEL_SIDE)
                if max_side > resize_target:
                    scale = resize_target / float(max_side)
                    new_size = (max(1, int(width * scale)), max(1, int(height * scale)))
                    image = image.resize(new_size, RESAMPLE)

                for quality in self.JPEG_QUALITY_STEPS:
                    buffer = io.BytesIO()
                    try:
                        image.save(buffer, format="JPEG", quality=quality, optimize=True)
                    except OSError:
                        image.save(buffer, format="JPEG", quality=quality)
                    data = buffer.getvalue()
                    if len(data) <= self.MAX_IMAGE_BYTES:
                        return data
        except Exception as exc:
            raise BaiduAPIError(f"压缩图片失败：{exc}")

        raise BaiduAPIError("图片压缩后仍超过百度接口限制，请手动压缩后再试")

    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    def classify_image(self, image_path: str) -> List[dict]:
        if not os.path.exists(image_path):
            raise FileNotFoundError(image_path)

        image_bytes = self._prepare_image_bytes(image_path)

        response = self._post_image(image_bytes)
        if "error_code" in response:
            raise BaiduAPIError(
                f"{response.get('error_code')}: {response.get('error_msg', '未知错误')}",
                response,
            )
        return response.get("result", [])

    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    def classify_files(
        self,
        file_paths: Iterable[str],
        progress_callback: Optional[Callable[[int, int, str], None]] = None,
    ) -> Tuple[Dict[str, List[str]], Dict[str, int]]:
        file_list = [os.fspath(path) for path in file_paths]
        total = len(file_list)
        classification: Dict[str, List[str]] = {}
        skipped = 0
        failed = 0
        errors: List[str] = []

        for idx, path in enumerate(file_list, start=1):
            filename = os.path.basename(path)
            if progress_callback:
                progress_callback(idx, total, filename)

            try:
                items = self.classify_image(path)
            except Exception as exc:
                failed += 1
                errors.append(f"{filename}: {exc}")
                continue

            folder = self._select_folder(items)
            if folder is None:
                skipped += 1
                continue

            classification.setdefault(folder, []).append(filename)

        stats = {"total": total, "skipped": skipped, "failed": failed, "errors": errors}
        return classification, stats

    # ------------------------------------------------------------------

    # ------------------------------------------------------------------
    def _select_folder(self, items: List[dict]) -> Optional[str]:
        if not items:
            return "未识别" if self.allow_unclassified else None

        for item in items:
            score = float(item.get("score", 0))
            keyword = str(item.get("keyword", "")).strip()
            if not keyword:
                continue
            if score >= self.min_score:
                return self._sanitize_folder_name(keyword)

        return "未识别" if self.allow_unclassified else None

    # ------------------------------------------------------------------
    def _sanitize_folder_name(self, name: str) -> str:
        cleaned = self.INVALID_CHARS.sub("_", name.strip())
        cleaned = cleaned.replace("\n", "_")
        if not cleaned:
            cleaned = "未命名分类"
        return cleaned[:80]

    # ------------------------------------------------------------------
    def _sanitize_folder_name(self, name: str) -> str:
        cleaned = self.INVALID_CHARS.sub("_", name.strip())
        cleaned = cleaned.replace("\n", "_")
        if not cleaned:
            cleaned = "未命名分类"
        return cleaned[:80]

